from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool

from src.ai.langchain.init_llm import get_llm


@tool
def compute_add(a: float, b: float) -> float:
    """计算两个数的和"""
    print("开始计算")
    return a + b


llm = get_llm()
tools = [compute_add]
llm_tools = llm.bind_tools(tools)

messages = [
    ("system", "你是一个解决日常问题的助手,需要准确的回答用户问题"),
    ("human", "{question}"),
]

prompt = ChatPromptTemplate.from_messages(messages)

chain = prompt | llm_tools

print(chain.invoke({"question": "调用工具计算3加4的和"}))


